import csv

from pymongo import MongoClient
from pymongo.errors import ConnectionFailure

# MongoDB连接配置
MONGO_HOST = '101.227.69.34'
MONGO_PORT = 2929
MONGO_USER = 'admin'
MONGO_PASS = "admin"
DB_NAME = 'hr'  # 指定数据库名
COLLECTION_NAME = 'employees'  # 指定集合名
try:
    # 创建MongoDB客户端连接
    client = MongoClient(
        host=MONGO_HOST,
        port=MONGO_PORT,
        username=MONGO_USER,
        password=MONGO_PASS,
        authSource='admin'  # 认证数据库
    )
    # 测试连接
    client.admin.command('ping')
    print("MongoDB连接成功！")
    # 可以添加更多测试代码，比如列出数据库
    print('数据库列表:', client.list_database_names())
    # 获取数据库和集合
    db = client[DB_NAME]
    collection = db[COLLECTION_NAME]
    # 读取CSV文件并导入数据
    with open("D:/zx/salary.csv", 'r', encoding='utf-8') as csvfile:
        csv_reader = csv.DictReader(csvfile)
        records = []
        for row in csv_reader:
            # 转换数据类型
            record = {
                'name': row['name'],
                'sex': row['sex'],
                'age': int(row['age']),  # 转换为整数
                'salary': int(row['salary'])  # 转换为整数
            }
            records.append(record)
        # 插入数据
        result = collection.insert_many(records)
        print(f"成功导入{len(result.inserted_ids)}条数据到{DB_NAME}.{COLLECTION_NAME}")
    # 显示导入的数据量
    print(f"当前集合文档数：{collection.count_documents({})}")
except ConnectionFailure as e:
    print(f"MongoDB连接失败：{e}")
except Exception as e:
    print(f"发生错误{e}")
finally:
    # 关闭连接
    if 'client' in locals():
        client.close()